import pretty_midi

def get_notes_with_duration(midi_file):
    # 加载 MIDI 文件
    midi_data = pretty_midi.PrettyMIDI(midi_file)
    notes_info = []

    # 遍历所有音轨
    for instrument in midi_data.instruments:
        # 只考虑具有音高的音轨
        if not instrument.is_drum:  # 过滤掉打击乐器
            for note in instrument.notes:
                note_info = {
                    'start_time': note.start,
                    'end_time': note.end,
                    'pitch': note.pitch,
                    'pitch_name': pretty_midi.note_number_to_name(note.pitch),  # 获取音符名称
                    'duration': note.end - note.start
                }
                notes_info.append(note_info)

    return notes_info

def classify_notes(notes_info):
    chords = []
    single_notes = []

    for i in range(len(notes_info)):
        current_note = notes_info[i]
        overlaps = []

        # 检查与其他音符的重叠
        for j in range(len(notes_info)):
            if i != j:  # 不与自身比较
                other_note = notes_info[j]
                # 判断时间是否重叠
                if (current_note['start_time'] < other_note['end_time'] and
                    current_note['end_time'] > other_note['start_time']):
                    overlaps.append(other_note)

        # 将当前音符分类为和弦或单音
        if overlaps:  # 如果有重叠，判断为和弦
            chord_pitches = {current_note['pitch_name']}
            chord_pitches.update(note['pitch_name'] for note in overlaps)
            chord_info = {
                'start_time': current_note['start_time'],
                'end_time': current_note['end_time'],
                'chord': '+'.join(chord_pitches),  # 组合音高
                'duration': current_note['duration']
            }
            chords.append(chord_info)
        else:  # 无重叠，判断为单音
            single_info = {
                'start_time': current_note['start_time'],
                'end_time': current_note['end_time'],
                'pitch': current_note['pitch_name'],
                'duration': current_note['duration']
            }
            single_notes.append(single_info)

    return chords, single_notes

def write_to_py_file(output_py_name='results_output.py', output_data=[], write_mode='num'):
    with open(output_py_name, "w", encoding="utf-8") as f:
        f.write("res_list = [\n")

        for item in output_data:
            if write_mode == 'num':
                formatted_item = [round(float(value), 5) for value in item]
                f.write(f"  {formatted_item},\n")
            else:
                f.write(f"  {item},\n")

        f.write("]\n")
    print(f"[INFO] Results saved to {output_py_name}.")

def main(midi_file):
    notes_info = get_notes_with_duration(midi_file)
    chords, single_notes = classify_notes(notes_info)

    # 导出为二维数组格式
    output_data = []
    for chord in chords:
        output_data.append([int(chord['start_time']*1000+1), int(chord['end_time']*1000-1), chord['chord']])

    for note in single_notes:
        output_data.append([int(note['start_time']*1000+1), int(note['end_time']*1000-1), note['pitch']])

    # 按起始时间排序
    output_data.sort(key=lambda x: x[0])  # 以起始时间排序

    # 输出为文件
    write_to_py_file(output_data=output_data, write_mode='string')

# 输入你的 MIDI 文件路径
midi_file_path = 'vio_basic_yinjie_3.mid'
main(midi_file_path)